package com.wht.ai.anythingLLM;

import lombok.SneakyThrows;
import org.noear.solon.ai.chat.ChatModel;
import org.noear.solon.ai.chat.ChatResponse;
import org.noear.solon.ai.chat.dialect.ChatDialectManager;
import org.noear.solon.rx.SimpleSubscriber;
import org.reactivestreams.Publisher;
;

import java.util.concurrent.CountDownLatch;

/**
 * @author by HaiTao.Wang on 2025/2/18.
 */
public class AnythingLLMTest {

    public static void main(String[] args) {

        ChatDialectManager.register(new AnythingLLMDialect());

//        chat("你是谁");

        chatStream("你是谁");

    }

    @SneakyThrows
    public static void chat(String prompt) {
        ChatModel chatModel = ChatModel.of("http://localhost:3001/api/v1/workspace/jt/chat")
                .headerSet("Authorization", "Bearer TY5KDCS-3EQ47FB-G1NAD4N-W4YXTD6")
                .provider("anythingLLM")
                .build();

        ChatResponse response = chatModel.prompt(prompt).call();
        System.err.println(response.getMessage());
    }

    @SneakyThrows
    public static void chatStream(String prompt) {
        ChatModel chatModel = ChatModel.of("http://localhost:3001/api/v1/workspace/jt/stream-chat")
                .headerSet("Authorization", "Bearer TY5KDCS-3EQ47FB-G1NAD4N-W4YXTD6")
                .provider("anythingLLM")
                .build();

        Publisher<ChatResponse> publisher = chatModel.prompt(prompt).stream();

        CountDownLatch latch = new CountDownLatch(1);
        publisher.subscribe(new SimpleSubscriber<ChatResponse>()
                .doOnNext(chatResponse -> {
                    System.err.println(chatResponse.getMessage());
                }).doOnComplete(() -> {
                    System.err.println("回答完毕");
                }));

        latch.await();

    }
}
